import numpy as np
import matplotlib.pyplot as plt

def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def perceptron(w, x, d, lr, is_binary=True):
    X = [1, x[0], x[1], x[2]]
    pred = sum([ww * xx for ww, xx in zip(w, X)])
    if is_binary:
        o = sigmoid(pred)
        do = o * (1 - o)
    else:
        o = np.sign(pred)
        do = 1.0
    loss = d - o
    w1 = [ww + lr * loss * do * xx for ww, xx in zip(w, X)]
    return loss**2, w1


def optimize(n, X, Y, lr, is_binary=True):
    w = np.array([0., 0., 0., 0.])
    losses = []
    for i in range(n):
        total_loss = 0
        for x, y in zip(X, Y):
            loss, w = perceptron(w, x, y, lr, is_binary)
            total_loss += loss
        losses.append(total_loss / 4)
    return losses


x1 = np.array([[1, 1, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1]])
y1 = np.array([1, 0, 0, 0])
x2 = np.array([[1, 1, 1], [1, 1, -1], [1, -1, 1], [-1, 1, 1]])
y2 = np.array([1, -1, -1, -1])

# compare lr
lrs = [10, 5, 2, 1, 0.5]
colors = ['green', 'orange', 'red', "yellow", "blue"]
plt.xlabel('iter')
plt.ylabel('loss')
n = 1000
for lr, c in zip(lrs, colors):
    losses = optimize(n, x1, y1, lr)
    plt.plot(range(n), losses, "-", color=c, label='learning_rate=%f' % (lr))
plt.legend()
plt.savefig("diff_lr.jpg")
plt.cla()

# compare binary
plt.xlabel('iter')
plt.ylabel('loss')
n = 500
lr = 0.2
losses = optimize(n, x1, y1, lr)
plt.plot(range(n), losses, "-", color="red", label='binary')
losses = optimize(n, x2, y2, lr, False)
plt.plot(range(n), losses, "-", color="blue", label='bipolar')
plt.legend()
plt.savefig("bin_bip.jpg")